from keras import layers, optimizers
from keras import models
from keras.preprocessing.image import ImageDataGenerator
import matplotlib.pyplot as plt

# 网络层构建
model = models.Sequential()
model.add(layers.Conv2D(32, (3, 3), activation='relu',input_shape=(150, 150, 3)))
model.add(layers.MaxPooling2D((2,2)))
model.add(layers.Conv2D(64, (3, 3), activation='relu'))
model.add(layers.MaxPooling2D((2,2)))
model.add(layers.Conv2D(128, (3, 3), activation='relu'))
model.add(layers.MaxPooling2D((2,2)))
model.add(layers.Conv2D(128, (3, 3), activation='relu'))
model.add(layers.MaxPooling2D((2,2)))

model.add(layers.Flatten())
model.add(layers.Dropout(0.5))
model.add(layers.Dense(512, activation='relu'))
model.add(layers.Dense(1, activation='sigmoid'))

# 查看特征图维度变化
# model.summary()

# 编译模型
model.compile(loss='binary_crossentropy',optimizer=optimizers.RMSprop(learning_rate=1e-4),metrics=['acc'])

# 处理数据

train_dir =  r'D:\Users\kk\PycharmProjects\pythonProject\pythonProject2\tt\demo02\cats_and_dogs_small\train'
validation_dir = r'D:\Users\kk\PycharmProjects\pythonProject\pythonProject2\tt\demo02\cats_and_dogs_small\validation'
#将所有图像乘以 1/255缩放
train_datagen = ImageDataGenerator(
    rescale=1./255,
    rotation_range=40,
    width_shift_range=0.2,
    height_shift_range=0.2,
    shear_range=0.2,
    zoom_range=0.2,
    horizontal_flip=True,)
test_datagen = ImageDataGenerator(rescale=1./255)

#训练数据生成器
train_generator = train_datagen.flow_from_directory(
    train_dir,                # 目标路径
    target_size=(150,150),    # 将所有图像大小调整为150×150
    batch_size=20,
    class_mode='binary')      # 二进制标签
#验证数据生成器
validation_generator = test_datagen.flow_from_directory(
    validation_dir,           # 目标路径
    target_size=(150,150),    # 将所有图像大小调整为150×150
    batch_size=20,
    class_mode='binary')      # 二进制标签
# 拟合模型
history = model.fit_generator(
    train_generator,
    steps_per_epoch=100,
    epochs=100,
    validation_data=validation_generator,
    validation_steps=50)
# 保存模型
model.save('cats_and_dogs_small_2.h5')
'''
绘制训练过程中的损失曲线和精度曲线
'''
acc = history.history['acc']
val_acc = history.history['val_acc']
loss = history.history['loss']
val_loss = history.history['val_loss']

epochs = range(1, len(acc) + 1)

plt.plot(epochs, acc, 'bo', label='Training acc')
plt.plot(epochs, val_acc, 'b', label='Validation acc')
plt.title('training and validation accuracy')
plt.legend()

plt.figure()

plt.plot(epochs, loss, 'bo', label='Training loss')
plt.plot(epochs, val_loss, 'b', label='Validation loss')
plt.title('training and validation loss')
plt.legend()

plt.show()